#include "znx_mem.h"

static size_t alignment = ZNX_MEM_4K_ALIGNMENT;

#if (ZNX_MEM_NO_POOL)

void
znx_mem_init(size_t water_level)
{
    // nothing
    (void)water_level;
}

void *
znx_mem_alloc(size_t size)
{
    znx_mem_item_head_t          *item;

    if (size > ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    item = malloc(sizeof(znx_mem_item_head_t) + size);
    item->magic = ZNX_MEM_ITEM_MAGIC;
    item->size = (uint32_t)size;
    item->next = NULL;

    return (void *)(item) + sizeof(znx_mem_item_head_t);
}

void
znx_mem_free(void *ptr)
{
    znx_mem_item_head_t          *item = NULL;
    item = ptr - sizeof(znx_mem_item_head_t);
    assert(item->magic == ZNX_MEM_ITEM_MAGIC);
    free(item);
}


void
znx_mem_status_report(char *buf, size_t buf_size)
{
    snprintf(buf, buf_size, "no pool\n");
}

void
znx_mem_clean()
{
    // nothing
}


void
znx_4kaligned_mem_init(size_t water_level)
{
    // nothing
    (void)water_level;
}


void *
znx_4kaligned_mem_alloc(size_t size)
{
    if (size >ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    size_t offset = (alignment - 1) + sizeof(znx_aligned_mem_item_head_t);
    znx_aligned_mem_item_head_t *item = malloc(size + offset);
    item->magic = ZNX_ALIGNED_MEM_ITEM_MAGIC;
    item->next = item->ptr = NULL;
    item->size = (uint32_t)size;
    void *p1 = (void *)item;
    void **p2 = (void**)(((size_t)p1 + offset) & ~(alignment - 1));
    p2[-1] = p1;
    return p2;
}


void
znx_4kaligned_mem_free(void *ptr)
{
    void                            *p1;
    znx_aligned_mem_item_head_t      *item;

    p1 = ((void **)ptr)[-1];
    item = (znx_aligned_mem_item_head_t *)p1;
    assert(item->magic == ZNX_ALIGNED_MEM_ITEM_MAGIC);
    free(item);
}


void
znx_4kaligned_mem_status_report(char *buf, size_t buf_size)
{
    snprintf(buf, buf_size, "no pool\n");
}


void
znx_4kaligned_mem_clean()
{
    // nothing
}


#else

static znx_mem_freelists_t       g_mem_freelists;

void
znx_mem_init(size_t water_level)
{
    int                 i;
    for (i = 0; i < ZNX_FREELIST_MODE_END; i++) {
        znx_spinlock_init(&g_mem_freelists.lists[i].lock);
        g_mem_freelists.lists[i].space_size
            = znx_mem_freelist_mode_to_size(i);
        g_mem_freelists.lists[i].water_level = water_level;
        g_mem_freelists.lists[i].count = 0;
        g_mem_freelists.lists[i].used = 0;
        g_mem_freelists.lists[i].head = (void *)(-1);
    }

    g_mem_freelists.inited = ZNX_TRUE;
}


static inline znx_mem_item_head_t *
znx_mem_base_alloc(size_t size)
{
    if (size > ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    znx_mem_item_head_t      *item;
    item = malloc(sizeof(znx_mem_item_head_t) + size);
    item->magic = ZNX_MEM_ITEM_MAGIC;
    item->size = (uint32_t)size;
    item->next = NULL;

    return item;
}

static inline void
znx_mem_base_free(znx_mem_item_head_t *item)
{
    free(item);
}


void *
znx_mem_alloc(size_t size)
{
    if (size > ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    int                     i;
    int                     mode = -1;
    size_t                  mode_size;

    znx_mem_item_head_t      *item = NULL;
    znx_mem_freelist_t       *freelist = NULL;

    for (i = 0; i < ZNX_FREELIST_MODE_END; i++) {
        mode_size = (size_t)znx_mem_freelist_mode_to_size(i);
        if (size <= mode_size) {
            size = mode_size;
            mode = i;
            break;
        }
    }

    if (mode != -1 && g_mem_freelists.inited) {
        freelist = &g_mem_freelists.lists[mode];
    }

    if (freelist != NULL) {
        znx_spinlock_lock(&freelist->lock);
        if (freelist->count > 0) {
            item = freelist->head;
            freelist->head = item->next;
            item->next = NULL;
            freelist->count--;
        }
        freelist->used++;
        znx_spinlock_unlock(&freelist->lock);
    }

    if (item == NULL) {
        item = znx_mem_base_alloc(size);
    }

    return (void *)(item) + sizeof(znx_mem_item_head_t);
}


void
znx_mem_free(void *ptr)
{
    znx_mem_item_head_t     *item = NULL;
    znx_mem_freelist_t      *freelist = NULL;
    int                     mode = -1;
    int                     i;

    item = ptr - sizeof(znx_mem_item_head_t);
    assert(item->magic == ZNX_MEM_ITEM_MAGIC);
    if (item->next != NULL) {
        ZNX_ASSERT_FATAL("double free or mem stampede");
    }

    for (i = 0; i < ZNX_FREELIST_MODE_END; i++) {
        if (item->size == znx_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode != -1 && g_mem_freelists.inited) {
        freelist = &g_mem_freelists.lists[mode];
    }

    if (freelist != NULL) {
        znx_spinlock_lock(&freelist->lock);
        if (freelist->count < freelist->water_level) {
            item->next = freelist->head;
            freelist->head = item;
            item = NULL;
            freelist->count++;
        }
        freelist->used--;
        znx_spinlock_unlock(&freelist->lock);
    }

    if (item != NULL) {
        znx_mem_base_free(item);
    }
}


void
znx_mem_status_report(char *buf, size_t buf_size)
{
    if (!g_mem_freelists.inited) {
        snprintf(buf, buf_size, "no inited\n");
        return;
    }

    int                 off = 0;
    int                 i;
    size_t              size;
    size_t              count;
    size_t              used;
    znx_mem_freelist_t  *freelist = NULL;
    for (i = 0; i < ZNX_FREELIST_MODE_END; i++) {
        freelist = &g_mem_freelists.lists[i];
        znx_spinlock_lock(&freelist->lock);
        size = freelist->space_size;
        count = freelist->count;
        used = freelist->used;
        znx_spinlock_unlock(&freelist->lock);

        off += snprintf(buf + off, buf_size - (size_t)off,
            "%ld mode, count: %ld, used: %ld\n", size, count, used);
    }
}


static void
znx_mem_freelist_clean(znx_mem_freelist_t *freelist)
{
    znx_mem_item_head_t     *item;

    znx_spinlock_lock(&freelist->lock);
    while (freelist->count > 0) {
        item = freelist->head;
        freelist->head = item->next;
        item->next = NULL;
        znx_mem_base_free(item);
        freelist->count--;
    }
    znx_spinlock_unlock(&freelist->lock);
}


void
znx_mem_clean()
{
    if (!g_mem_freelists.inited) {
        return;
    }

    int                 i;

    for (i = 0; i < ZNX_FREELIST_MODE_END; i++) {
        znx_mem_freelist_clean(&g_mem_freelists.lists[i]);
    }
}


static znx_4kaligned_mem_freelists_t        g_4kaligned_mem_freelists;

void
znx_4kaligned_mem_init(size_t water_level)
{
    int                         i;
    for (i = 0; i < MEM_4KALIGNED_FREELIST_MODE_END; i++) {
        znx_spinlock_init(&g_4kaligned_mem_freelists.lists[i].lock);
        g_4kaligned_mem_freelists.lists[i].space_size
            = znx_4kaligned_mem_freelist_mode_to_size(i);
        g_4kaligned_mem_freelists.lists[i].water_level = water_level;
        g_4kaligned_mem_freelists.lists[i].count = 0;
        g_4kaligned_mem_freelists.lists[i].used = 0;
        g_4kaligned_mem_freelists.lists[i].head = (void *)(-1);
    }

    g_4kaligned_mem_freelists.inited = ZNX_TRUE;
}


static inline znx_aligned_mem_item_head_t *
znx_4kaligned_mem_base_alloc(size_t size)
{
    if (size > ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    size_t offset = (alignment - 1) + sizeof(znx_aligned_mem_item_head_t);
    znx_aligned_mem_item_head_t *item = malloc(size + offset);
    item->magic = ZNX_ALIGNED_MEM_ITEM_MAGIC;
    item->next = item->ptr = NULL;
    item->size = (uint32_t)size;
    void *p1 = (void *)item;
    void **p2 = (void**)(((size_t)p1 + offset) & ~(alignment - 1));
    p2[-1] = p1;
    return item;
}


static inline void
znx_4kaligned_mem_base_free(znx_aligned_mem_item_head_t *item)
{
    free(item);
}


void *
znx_4kaligned_mem_alloc(size_t size)
{
    if (size > ZNX_MEM_MAX_ALLOC_SIZE) {
        ZNX_ASSERT_FATAL("to big alloc size");
    }

    znx_mem_freelist_t           *freelist = NULL;
    znx_aligned_mem_item_head_t  *item = NULL;
    int                         i;
    int                         mode = -1;
    size_t                      mode_size;

    for (i = 0; i < MEM_4KALIGNED_FREELIST_MODE_END; i++) {
        mode_size = (size_t)znx_4kaligned_mem_freelist_mode_to_size(i);
        if (size <= mode_size) {
            size = mode_size;
            mode = i;
            break;
        }
    }

    if (mode != -1 && g_4kaligned_mem_freelists.inited) {
        freelist = &g_4kaligned_mem_freelists.lists[mode];
    }

    if (freelist != NULL) {
        znx_spinlock_lock(&freelist->lock);
        if (freelist->count > 0) {
            item = freelist->head;
            freelist->head = item->next;
            item->next = NULL;
            freelist->count--;
        }
        freelist->used++;
        znx_spinlock_unlock(&freelist->lock);
    }

    if (item == NULL) {
        item = znx_4kaligned_mem_base_alloc(size);
    }

    return ZNX_4KALIGNED_MEM_ITEM_PTR(item);
}


void
znx_4kaligned_mem_free(void *ptr)
{
    void                            *p1;
    znx_aligned_mem_item_head_t     *item;
    znx_mem_freelist_t              *freelist = NULL;
    int                             mode = -1;
    int                             i;

    p1 = ((void **)ptr)[-1];
    item = (znx_aligned_mem_item_head_t *)p1;
    assert(item->magic == ZNX_ALIGNED_MEM_ITEM_MAGIC);
    if (item->next != NULL) {
        ZNX_ASSERT_FATAL("double free or mem stampede");
    }

    for (i = 0; i < MEM_4KALIGNED_FREELIST_MODE_END; i++) {
        if (item->size == znx_4kaligned_mem_freelist_mode_to_size(i)) {
            mode = i;
            break;
        }
    }

    if (mode != -1 && g_4kaligned_mem_freelists.inited) {
        freelist = &g_4kaligned_mem_freelists.lists[mode];
    }

    if (freelist != NULL) {
        znx_spinlock_lock(&freelist->lock);
        if (freelist->count < freelist->water_level) {
            item->next = freelist->head;
            freelist->head = item;
            item = NULL;
            freelist->count++;
        }
        freelist->used--;
        znx_spinlock_unlock(&freelist->lock);
    }

    if (item != NULL) {
        znx_4kaligned_mem_base_free(item);
    }
}


void
znx_4kaligned_mem_status_report(char *buf, size_t buf_size)
{
    if (!g_4kaligned_mem_freelists.inited) {
        snprintf(buf, buf_size, "no inited\n");
        return;
    }

    int                     off = 0;
    int                     i;
    size_t                  size;
    size_t                  count;
    size_t                  used;
    znx_mem_freelist_t      *freelist = NULL;
    for (i = 0; i < MEM_4KALIGNED_FREELIST_MODE_END; i++) {
        freelist = &g_4kaligned_mem_freelists.lists[i];
        znx_spinlock_lock(&freelist->lock);
        size = freelist->space_size;
        count = freelist->count;
        used = freelist->used;
        znx_spinlock_unlock(&freelist->lock);

        off += snprintf(buf + off, buf_size - (size_t)off,
            "%ld aligned mode, count: %ld, used: %ld\n", size, count, used);
    }
}


static void
znx_4kaligned_mem_freelist_clean(znx_mem_freelist_t *freelist)
{
    znx_aligned_mem_item_head_t     *item;

    znx_spinlock_lock(&freelist->lock);
    while (freelist->count > 0) {
        item = freelist->head;
        freelist->head = item->next;
        item->next = NULL;
        znx_4kaligned_mem_base_free(item);
        freelist->count--;
    }
    znx_spinlock_unlock(&freelist->lock);
}


void
znx_4kaligned_mem_clean()
{
    if (!g_4kaligned_mem_freelists.inited) {
        return;
    }

    int                             i;

    for (i = 0; i < MEM_4KALIGNED_FREELIST_MODE_END; i++) {
        znx_4kaligned_mem_freelist_clean(&g_4kaligned_mem_freelists.lists[i]);
    }
}

#endif
